# 示例：create_react_agent_memory.py

"""
REACT + Checkpoint 简介：
REACT 是 LangGraph 提供的一种智能代理设计模式，它结合了：
1. Reasoning（推理）：代理可以根据输入做逻辑判断。
2. Acting（执行）：代理可以调用工具（比如天气查询）。
3. Checkingpoint（检查点/记忆）：通过 checkpoint 保存代理状态和上下文，
   使得每次调用代理时都能保留之前的聊天历史或计算结果，从而支持连续对话或复杂任务。

简单理解：
- REACT 代理就像一个智能助理，既能思考又能行动。
- Checkpoint 就像笔记本，帮助代理记住之前的内容。
"""

# ----------------- 模型初始化 -----------------
import os
from dotenv import load_dotenv
from typing import Literal

# langchain 工具装饰器，用于把函数包装成可供代理调用的工具
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI

# 自动加载 .env 文件中的 API key 等敏感信息
load_dotenv()

# 初始化 LLM 模型
# - model: 指定使用的模型
# - streaming: 是否开启流式输出
# - temperature: 控制生成文本的随机性，越低越确定
llm = ChatOpenAI(
    api_key=os.getenv("MODELSCOPE_API_KEY"),
    openai_api_base=os.getenv("MODELSCOPE_API_BASE"),
    model="Qwen/Qwen3-32B",
    streaming=True,
    temperature=0.7
)

# ----------------- 定义工具 -----------------
@tool
def get_weather(city: Literal["nyc", "sf"]):
    """
    工具函数：根据城市返回预定义天气信息
    REACT 代理可以调用此工具来“行动”
    """
    if city == "nyc":
        return "It might be cloudy in nyc"
    elif city == "sf":
        return "It's always sunny in sf"
    else:
        raise AssertionError("Unknown city")

# 将工具放入列表，方便代理使用
tools = [get_weather]

# ----------------- 配置 Checkpoint -----------------
from langgraph.checkpoint.memory import MemorySaver

# 初始化 MemorySaver，用于保存聊天记忆
memory = MemorySaver()

# ----------------- 创建 REACT 代理 -----------------
from langgraph.prebuilt import create_react_agent

# 创建 REACT 代理：
# - llm: 使用的语言模型
# - tools: 可调用的工具
# - checkpointer: 负责保存/恢复代理状态的 checkpoint
graph = create_react_agent(llm, tools=tools, checkpointer=memory)

# ----------------- 定义打印函数 -----------------
def print_stream(stream):
    """
    用于打印流式输出
    - stream: graph.stream() 返回的迭代器
    - 支持 message 是 tuple 或对象
    """
    for s in stream:
        message = s["messages"][-1]
        if isinstance(message, tuple):
            print(message)
        else:
            # pretty_print() 会更清晰地输出代理生成的内容
            message.pretty_print()

# ----------------- 执行示例 -----------------
# 配置参数，可以设置 thread_id 保持同一个对话线程
config = {"configurable": {"thread_id": "1"}}

# 输入消息
inputs = {"messages": [("user", "What's the weather in NYC?")]}

# 打印流输出
print_stream(graph.stream(inputs, config=config, stream_mode="values"))
